import requests
import json
import re
import base64
from datetime import datetime, timedelta

# ========== DeepSeek 配置 ==========
API_KEY = "sk-" 
API_URL = "https://api.deepseek.com/v1/chat/completions"
HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}


def encode_image_to_base64(filepath):
    try:
        with open(filepath, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")
    except FileNotFoundError:
        print(f"[错误] 图片路径不存在: {filepath}")
        return ""

def millis_to_iso(millis):
    return datetime.utcfromtimestamp(int(millis) / 1000).isoformat() + "Z"

def extract_json_from_response(response_text):
    match = re.search(r"```json\s*(.*?)\s*```", response_text, re.DOTALL)
    json_str = match.group(1) if match else response_text
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        print("无法解析 JSON：", e)
        return None

def filter_and_format_messages(raw_messages):
    filtered = []
    for msg in raw_messages:
        if msg["msg_type"] == "text":
            text = msg["message_content"]["text"]
            if not any(x in text for x in ["收到", "谢谢", "了解了"]):
                filtered.append({
                    "msg_type": "text",
                    "message_content": {
                        "text": text
                    },
                    "create_time": millis_to_iso(msg.get("create_time", ""))
                })
        elif msg["msg_type"] == "image":
            filtered.append({
                "msg_type": "image",
                "message_content": {
                    "image_key": msg["message_content"]["image_key"],
                    "base64": msg.get("base64", "")
                },
                "create_time": millis_to_iso(msg.get("create_time", ""))
            })
    return filtered

def cluster_and_classify(messages):
    prompt_lines = [
        f"- [{msg['create_time']}] [TEXT] {msg['message_content']['text']}"
        for msg in messages if msg["msg_type"] == "text"
    ]

    prompt = f"""
<context>
你是一个建筑工地巡检助手。在巡检过程中，你需要将信息聚类、分类，并基于每个聚类生成专业的问题描述，并提取关键信息。
</context>

<task_rules>

对于以下巡检留言列表，请执行以下任务：

1. 将留言按照内容相关性进行聚类，每个聚类形成一个事件卡；
2. 对每个事件卡进行分类，类别为以下五类之一：
    - 巡检问题（现场隐患、风险点）
    - 验收
    - 旁站
    - 闭环
    - 其他
3. 对于被分类为“问题记录”的事件卡，还需要：
    - 生成一段专业、规范的“问题描述”（要求简洁清晰）
    - 尝试提取：
        - 施工部位 (Location)
        - 责任单位 (Responsible Party)
    - 如果施工部位或责任单位缺失，请在提取内容中标记为 "[部位信息缺失]" 或 "[责任单位信息缺失]"

</task_rules>

<output_rules>
请以以下 JSON 数组格式返回所有事件卡，每个事件卡包含：

{{
  "event_id": int,             // 事件编号
  "summary": string,           // 事件总结
  "messages": [ ... ],         // 聚类包含的原始留言 (保留 msg_type, message_content, create_time)
  "class": string,             // 事件分类，若是巡检问题则填"巡检问题"，否则用其他分类
  "issue_description": string, // (仅巡检问题填写) AI生成的问题描述
  "location": string,          // (仅巡检问题填写) 施工部位或 "[部位信息缺失]"
  "responsible_party": string  // (仅巡检问题填写) 责任单位或 "[责任单位信息缺失]"
}}
</output_rules>

以下是巡检留言信息：
{chr(10).join(prompt_lines)}
"""
    
    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "你是一个建筑巡检信息聚类、分类与关键信息提取助手。"},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload)
        response.raise_for_status()
        return extract_json_from_response(response.json()["choices"][0]["message"]["content"])
    except requests.exceptions.RequestException as e:
        print("请求失败:", e)
        return None

def assign_images_to_events(events, images, time_window_minutes=10):
    for event in events:
        event_time_range = [
            datetime.fromisoformat(msg["create_time"].replace("Z", "+00:00"))
            for msg in event["messages"]
        ]
        min_time = min(event_time_range)
        max_time = max(event_time_range)
        window_start = min_time - timedelta(minutes=time_window_minutes)
        window_end = max_time + timedelta(minutes=time_window_minutes)
        #print(window_start, window_end)

        candidate_imgs = []
        for img in images:
            try:
                img_time = datetime.fromisoformat(img["create_time"].replace("Z", "+00:00"))
                if window_start <= img_time <= window_end:
                    candidate_imgs.append({
                        "base64": img["message_content"]["base64"],
                        "filename": img["message_content"]["image_key"],
                        "create_time": img_time
                    })
            except:
                continue

        event["candidate_images"] = candidate_imgs
    return events

def print_events(events):
    print("\n最终事件卡：")
    
    if not events:
        print("没有任何事件卡。")
        return
    
    for e in events:
        print("\n")
        print(f"事件 {e.get('event_id', '[无ID]')}: [{e.get('class', '[无类别]')}] - {e.get('summary', '[无总结]')}")
        print("  Messages:")

        if not e.get('messages'):
            print("    (无留言信息)")
        else:
            for msg in e['messages']:
                msg_type = msg.get('msg_type', '[未知]')
                mc = msg.get('message_content', '')
                timestamp = msg.get('create_time', '[无时间]')

                # 判断 message_content 是 dict 还是 str
                if isinstance(mc, dict):
                    content = mc.get("text") or mc.get("image_key") or "[空内容]"
                elif isinstance(mc, str):
                    content = mc
                else:
                    content = "[无法识别的内容格式]"

                print(f"    - [{msg_type}] {timestamp} | {content}")

        candidate_images = e.get('candidate_images', [])
        print(f"  Candidate Images: {len(candidate_images)}")
        
        # 如果是问题记录，打印更详细内容
        if e.get("class") == "问题记录":
            print(f"  问题描述: {e.get('issue_description', '[无描述]')}")
            print(f"  施工部位: {e.get('location', '[无施工部位]')}")
            print(f"  责任单位: {e.get('responsible_party', '[无责任单位]')}")


def main():
    img2 = encode_image_to_base64("IMG_2.jpg")
    img3 = encode_image_to_base64("image03.jpg")

    raw_messages = [
        {
            "message_id": "om_3",
            "msg_type": "text",
            "create_time": "1711933800000",  # 2024-03-31 22:30:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "在3号楼东南角发现脚手架松动，有坠落风险"}
        },
        {
            "message_id": "om_5",
            "msg_type": "text",
            "create_time": "1711933870000",  # 2024-03-31 22:31:10 UTC
            "sender_id": "user1",
            "message_content": {"text": "责任人张三已通知整改"}
        },
        {
            "message_id": "om_6",
            "msg_type": "text",
            "create_time": "1711934100000",  # 2024-03-31 22:35:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "收到，谢谢"}
        },
        {
            "message_id": "om_7",
            "msg_type": "text",
            "create_time": "1711940400000",  # 2024-03-31 23:00:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "2号楼发现电缆裸露，存在触电危险"}
        },
        {
            "message_id": "om_8",
            "msg_type": "image",
            "create_time": "1711940460000",  # 2024-03-31 23:01:00 UTC
            "sender_id": "user2",
            "message_content": {"image_key": "IMG_2.jpg"},
            "base64": img2
        },
        {
            "message_id": "om_9",
            "msg_type": "text",
            "create_time": "1711940700000",  # 2024-03-31 23:05:00 UTC
            "sender_id": "user2",
            "message_content": {"text": "监理回复：了解了"}
        },
        {
            "message_id": "om_extra_07",
            "msg_type": "text",
            "create_time": "1711952400000",  # 2024-04-01 02:00:00 UTC
            "sender_id": "user5",
            "message_content": {"text": "今日上午9点，3号楼主体钢筋绑扎作业开始"}
        },
        {
            "message_id": "om_extra_08",
            "msg_type": "text",
            "create_time": "1711953300000",  # 2024-04-01 02:15:00 UTC
            "sender_id": "user5",
            "message_content": {"text": "监理张工已到场确认施工规范"}
        },
        {
            "message_id": "om_extra_05",
            "msg_type": "text",
            "create_time": "1712028300000",  # 2024-04-02 08:45:00 UTC
            "sender_id": "user4",
            "message_content": {"text": "5号楼地下室防水层施工于昨日完成，今日进行初步验收"}
        },
        {
            "message_id": "om_extra_06",
            "msg_type": "text",
            "create_time": "1712028600000",  # 2024-04-02 08:50:00 UTC
            "sender_id": "user4",
            "message_content": {"text": "验收人员已到场，准备进行验收"}
        },
        {
            "message_id": "om_a",
            "msg_type": "image",
            "create_time": "1712029900000",  
            "sender_id": "user4",
            "message_content": {"image_key": "image03.jpg"},
            "base64": img3
        }
    ]

    filtered = filter_and_format_messages(raw_messages)
    text_msgs = [m for m in filtered if m["msg_type"] == "text"]
    image_msgs = [m for m in filtered if m["msg_type"] == "image"]

    print("\n所有图像消息:")
    for im in image_msgs:
        print(f"- {im['message_content']['image_key']} @ {im['create_time']}")



    events = cluster_and_classify(text_msgs)
    if events is None:
        print("聚类和分类失败")
        return

    events = assign_images_to_events(events, image_msgs)

    print_events(events)


if __name__ == "__main__":
    main()
